import os
import sys
import argparse
import hashlib
import json
import zlib
from pathlib import Path
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
import secrets
import base64
import time

# 加密相关函数
def generate_key():
    """生成随机加密密钥"""
    return secrets.token_bytes(32)  # AES-256需要32字节密钥

def encrypt_data(data, key):
    """使用AES-CBC加密数据"""
    # 生成随机初始化向量
    iv = secrets.token_bytes(16)
    # 创建加密器
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    
    # 填充数据到16字节倍数
    padder = padding.PKCS7(128).padder()  # AES块大小为128位
    padded_data = padder.update(data) + padder.finalize()
    
    # 加密数据
    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
    
    # 返回IV和加密后的数据
    return iv + encrypted_data

def decrypt_data(encrypted_data, key):
    """使用AES-CBC解密数据"""
    # 提取IV
    iv = encrypted_data[:16]
    encrypted_data = encrypted_data[16:]
    
    # 创建解密器
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    
    # 解密数据
    decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
    
    # 去除填充
    unpadder = padding.PKCS7(128).unpadder()
    try:
        decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
        return decrypted_data
    except ValueError as e:
        print(f"Padding error: {e}")
        # 返回未填充的数据（用于调试）
        return decrypted_padded_data

def embed_password(password):
    """将密码转换为可嵌入的格式"""
    # 使用URL安全的Base64编码，确保所有字符都是URL安全的
    encoded = base64.urlsafe_b64encode(password.encode('utf-8')).decode('ascii')
    # 添加版本标记，以便将来可以支持不同的编码方式
    return f"v1:{encoded}"

def extract_password(embedded_password):
    """从嵌入格式提取密码"""
    # 检查版本标记
    if embedded_password.startswith("v1:"):
        embedded_password = embedded_password[3:]
    
    try:
        # 使用URL安全的Base64解码
        return base64.urlsafe_b64decode(embedded_password.encode('ascii')).decode('utf-8')
    except Exception as e:
        print(f"Error decoding password: {e}")
        return None

# 进度条功能
def print_progress(progress, total, prefix='', suffix='', bar_length=50):
    """打印进度条"""
    if total == 0:
        percent = 100
    else:
        percent = int(100 * (progress / total))
    filled_length = int(bar_length * progress // total) if total > 0 else bar_length
    bar = '█' * filled_length + '-' * (bar_length - filled_length)
    sys.stdout.write(f'\r{prefix} |{bar}| {percent}% {suffix}')
    if progress >= total:
        sys.stdout.write('\n')
    sys.stdout.flush()

def get_total_size(sources):
    """计算所有源文件的总大小"""
    total_size = 0
    for source in sources:
        source_path = Path(source)
        if not source_path.exists():
            continue
        
        if source_path.is_file():
            total_size += source_path.stat().st_size
        elif source_path.is_dir():
            for root, _, files in os.walk(source_path):
                for file in files:
                    file_path = Path(root) / file
                    total_size += file_path.stat().st_size
    
    return total_size

def create_package(sources, output_file, password=None):
    """创建打包文件，支持加密"""
    package_data = {
        "files": [],
        "encrypted": False
    }
    
    
    # 临时文件用于写入
    temp_file = output_file + '.tmp'
    
    # 计算总大小用于进度条
    total_size = get_total_size(sources)
    processed_size = 0
    
    try:
        with open(temp_file, 'wb') as pkg_file:
            # 写入头部标识
            pkg_file.write(b'PKG1')
            
            # 处理加密
            if password:
                # 生成盐值
                salt = os.urandom(16)
                
                # 派生密钥
                key = hashlib.scrypt(
                    password.encode('utf-8'), 
                    salt=salt, 
                    n=2**14, r=8, p=1, 
                    dklen=32
                )
                
                package_data["encrypted"] = True
                package_data["salt"] = salt.hex()  # 保存盐值
                
                # 写入加密标记和盐值
                pkg_file.write(b'ENCR')
                pkg_file.write(salt)  # 写入盐值
                
                # 嵌入密码（以隐藏形式）
                embedded_password = embed_password(password)
                password_bytes = embedded_password.encode('utf-8')
                pkg_file.write(len(password_bytes).to_bytes(2, 'big'))
                pkg_file.write(password_bytes)
                
                # 记录密码哈希，用于验证
                password_hash = hashlib.sha256(password.encode('utf-8')).hexdigest()
                package_data["password_hash"] = password_hash
            else:
                # 写入未加密标记
                pkg_file.write(b'PLAI')
            
            # 处理每个源文件或目录
            for source in sources:
                source_path = Path(source)
                if not source_path.exists():
                    print(f"Error: Source {source} does not exist.")
                    continue
                
                if source_path.is_file():
                    # 处理单个文件
                    file_info = {
                        "path": str(source_path),
                        "offset": pkg_file.tell(),
                    }
                    
                    # 读取文件内容
                    file_size = source_path.stat().st_size
                    with open(source_path, 'rb') as f:
                        file_data = f.read()
                    
                    # 更新进度条
                    processed_size += file_size
                    print_progress(processed_size, total_size, prefix='打包中:', suffix=f'处理 {source_path.name}')
                    
                    # 计算哈希（在加密前）
                    file_info["sha256"] = hashlib.sha256(file_data).hexdigest()
                    
                    # 加密文件内容（如果需要）
                    if password:
                        encrypted_data = encrypt_data(file_data, key)
                        pkg_file.write(encrypted_data)
                        file_info["size"] = len(encrypted_data)
                        file_info["encrypted"] = True
                    else:
                        pkg_file.write(file_data)
                        file_info["size"] = len(file_data)
                        file_info["encrypted"] = False
                    
                    package_data["files"].append(file_info)
                
                elif source_path.is_dir():
                    # 处理目录（递归）
                    dir_files = []
                    for root, _, files in os.walk(source_path):
                        for file in files:
                            file_path = Path(root) / file
                            dir_files.append(file_path)
                    
                    for i, file_path in enumerate(dir_files):
                        # 使用相对路径基于源目录而非当前工作目录
                        rel_path = file_path.relative_to(source_path.parent)
                        
                        file_info = {
                            "path": str(rel_path),
                            "offset": pkg_file.tell(),
                        }
                        
                        # 读取文件内容
                        file_size = file_path.stat().st_size
                        with open(file_path, 'rb') as f:
                            file_data = f.read()
                        
                        # 更新进度条
                        processed_size += file_size
                        print_progress(processed_size, total_size, prefix='打包中:', suffix=f'处理 {rel_path} ({i+1}/{len(dir_files)})')
                        
                        # 计算哈希（在加密前）
                        file_info["sha256"] = hashlib.sha256(file_data).hexdigest()
                        
                        # 加密文件内容（如果需要）
                        if password:
                            encrypted_data = encrypt_data(file_data, key)
                            pkg_file.write(encrypted_data)
                            file_info["size"] = len(encrypted_data)
                            file_info["encrypted"] = True
                        else:
                            pkg_file.write(file_data)
                            file_info["size"] = len(file_data)
                            file_info["encrypted"] = False
                        
                        package_data["files"].append(file_info)
            
            # 写入元数据位置标记
            metadata_offset = pkg_file.tell()
            pkg_file.write(b'META')
            
            # 写入元数据大小（占位）
            metadata_size_pos = pkg_file.tell()
            pkg_file.write(b'\x00\x00\x00\x00')
            
            # 写入元数据
            metadata = json.dumps(package_data).encode('utf-8')
            compressed_metadata = zlib.compress(metadata)
            
            # 加密元数据（如果需要）
            if password:
                encrypted_metadata = encrypt_data(compressed_metadata, key)
                pkg_file.write(encrypted_metadata)
                metadata_size = len(encrypted_metadata)
                package_data["meta_encrypted"] = True
            else:
                pkg_file.write(compressed_metadata)
                metadata_size = len(compressed_metadata)
                package_data["meta_encrypted"] = False
            
            # 返回并写入实际的元数据大小
            current_pos = pkg_file.tell()
            
            pkg_file.seek(metadata_size_pos)
            pkg_file.write(metadata_size.to_bytes(4, 'big'))
            pkg_file.seek(current_pos)
            
            # 写入文件结束标记
            pkg_file.write(b'END\x00')
        
        # 替换为最终文件
        os.replace(temp_file, output_file)
        print(f"Package created: {output_file}")
        
    except Exception as e:
        print(f"Error creating package: {e}")
        # 清理临时文件
        if os.path.exists(temp_file):
            os.remove(temp_file)

def extract_package(package_file, output_dir='.', password=None):
    """解包文件到指定目录，支持解密"""
    # 确保输出目录存在
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    try:
        with open(package_file, 'rb') as pkg_file:
            # 验证文件头部
            header = pkg_file.read(4)
            if header != b'PKG1':
                print("Error: Not a valid package file.")
                return
            
            # 检查加密状态
            encryption_flag = pkg_file.read(4)
            is_encrypted = encryption_flag == b'ENCR'
            
            if is_encrypted:
                # 读取盐值
                salt = pkg_file.read(16)
                if len(salt) != 16:
                    print("Error: Invalid salt in encrypted package.")
                    return
                
                # 读取嵌入的密码
                password_len_bytes = pkg_file.read(2)
                password_len = int.from_bytes(password_len_bytes, 'big')
                embedded_password = pkg_file.read(password_len).decode('utf-8')
                password = extract_password(embedded_password)
                
                if not password:
                    print("Error: Failed to extract password from package.")
                    return
                
                # 派生密钥
                key = hashlib.scrypt(
                    password.encode('utf-8'), 
                    salt=salt, 
                    n=2**14, r=8, p=1, 
                    dklen=32
                )
                
                print(f"Decrypting package with password: {password[:3]}...")
            else:
                if password:
                    print("Warning: Password provided, but package is not encrypted.")
            
            # 查找元数据位置
            file_size = os.path.getsize(package_file)
            
            # 寻找文件结束标记 "END\x00"
            end_marker = b'END\x00'
            end_marker_size = len(end_marker)
            
            end_position = -1
            # 从文件末尾向前搜索结束标记
            for i in range(file_size - end_marker_size, -1, -1):
                pkg_file.seek(i)
                if pkg_file.read(end_marker_size) == end_marker:
                    end_position = i
                    break
            
            if end_position == -1:
                print("Error: Could not find file end marker.")
                return
            
            # 从结束标记向前查找META标记
            meta_marker = b'META'
            meta_marker_size = len(meta_marker)
            meta_position = -1
            
            # 从结束标记位置向前查找META标记
            for j in range(end_position - meta_marker_size, -1, -1):
                pkg_file.seek(j)
                if pkg_file.read(meta_marker_size) == meta_marker:
                    meta_position = j
                    break
            
            if meta_position == -1:
                print("Error: Could not find metadata marker.")
                return
            
            # 读取元数据大小
            metadata_size_pos = meta_position + meta_marker_size
            pkg_file.seek(metadata_size_pos)
            meta_size_bytes = pkg_file.read(4)
            meta_size = int.from_bytes(meta_size_bytes, 'big')
            
            # 读取元数据
            metadata_pos = metadata_size_pos + 4
            pkg_file.seek(metadata_pos)
            metadata_data = pkg_file.read(meta_size)
            
            # 解密元数据（如果需要）
            if is_encrypted:
                try:
                    decrypted_metadata = decrypt_data(metadata_data, key)
                    compressed_metadata = decrypted_metadata
                    
                    # 尝试解压元数据，验证解密是否成功
                    metadata = json.loads(zlib.decompress(compressed_metadata))
                    
                    # 验证密码哈希（如果存在）
                    if 'password_hash' in metadata:
                        actual_hash = hashlib.sha256(password.encode('utf-8')).hexdigest()
                        if actual_hash != metadata['password_hash']:
                            print("Warning: Password hash mismatch. Decryption may fail.")
                    
                except Exception as e:
                    print(f"Error decrypting metadata: {e}. Incorrect password or corrupted data.")
                    # 输出一些调试信息
                    print(f"Metadata size: {meta_size} bytes")
                    print(f"Extracted password: {password}")
                    print(f"Salt: {salt.hex()}")
                    return
            else:
                compressed_metadata = metadata_data
            
            # 解压元数据
            metadata = json.loads(zlib.decompress(compressed_metadata))
            
            # 计算总大小用于进度条
            files = metadata.get('files', [])
            total_size = sum(file.get('size', 0) for file in files)
            processed_size = 0
            
            # 提取文件
            for i, file_info in enumerate(files):
                file_path = Path(file_info['path'])
                offset = file_info['offset']
                size = file_info.get('size', 0)
                expected_sha256 = file_info['sha256']
                file_encrypted = file_info.get('encrypted', False)
                
                # 构建完整的输出路径
                output_file_path = output_path / file_path
                
                # 更新进度条
                print_progress(processed_size, total_size, prefix='解包中:', suffix=f'处理 {file_path} ({i+1}/{len(files)})')
                
                # 确保加密状态一致
                if file_encrypted != is_encrypted:
                    print(f"Warning: Inconsistent encryption state for {file_path}. Skipping.")
                    continue
                
                # 创建目录结构（如果不存在）
                output_file_path.parent.mkdir(parents=True, exist_ok=True)
                
                # 读取文件数据
                pkg_file.seek(offset)
                file_data = pkg_file.read(size)
                
                # 解密文件（如果需要）
                if file_encrypted:
                    try:
                        decrypted_data = decrypt_data(file_data, key)
                        file_data = decrypted_data
                    except Exception as e:
                        print(f"Error decrypting {file_path}: {e}. Skipping.")
                        continue
                
                # 更新进度（包括解密后的数据大小）
                processed_size += size
                
                # 验证哈希
                actual_sha256 = hashlib.sha256(file_data).hexdigest()
                if actual_sha256 != expected_sha256:
                    print(f"Warning: Hash mismatch for {file_path}. Possible corruption.")
                
                # 写入文件
                with open(output_file_path, 'wb') as f:
                    f.write(file_data)
            
            # 确保进度条显示100%
            print_progress(total_size, total_size, prefix='解包中:', suffix='完成')
            print(f"Package extracted successfully to {output_dir}")
            
    except Exception as e:
        print(f"Error extracting package: {e}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='文件打包解包工具')
    
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-p', '--pack', nargs=2, metavar=('SOURCES', 'OUTPUT'),
                       help='打包文件或目录，格式: -p "文件1,文件2,目录1" 输出文件名')
    group.add_argument('-d', '--depack', nargs=1, metavar=('PACKAGE'),
                       help='解包文件，格式: -d 包文件名.pkg')
    
    parser.add_argument('-e', '--encrypt', action='store_true',
                       help='启用加密功能（打包时需要）')
    parser.add_argument('-P', '--password', help='加密密码（仅打包时需要）')
    parser.add_argument('-o', '--output', default='.', help='解包输出目录（默认当前目录）')
    
    args = parser.parse_args()
    
    if args.pack:
        sources_str, output_file = args.pack
        sources = [s.strip() for s in sources_str.split(',')]
        
        if args.encrypt:
            if not args.password:
                print("Error: 启用加密时必须提供密码 (-P)")
                sys.exit(1)
            create_package(sources, output_file, args.password)
        else:
            create_package(sources, output_file)
    elif args.depack:
        package_file = args.depack[0]
        # 解包时不需要密码参数，自动从包中读取
        extract_package(package_file, args.output)

if __name__ == "__main__":
    main()    